#include "DeadCodeEli.h"
#include "RDominateTree.h"
#include "IRStmtBuilder.h"

static bool IsCritical(Instruction* inst);

void DeadCodeEli::execute()
{
    module->set_print_name();
    init_rdf();
    for(auto fun: module->get_functions())
    {
        if (fun->get_basic_blocks().empty()) continue;
        func_ = fun;
        init_worklist();
        mark();
        sweep();
    }
}

void DeadCodeEli::init_rdf()
{
    RDominateTree rdt(module);
    rdt.execute();
}

void DeadCodeEli::init_worklist()
{
    worklist.clear();
    marked.clear();
    for (auto bb : func_->get_basic_blocks())
        for (auto inst : bb->get_instructions())
        {
            if (IsCritical(inst))
            {
                worklist[inst] = true;
                marked[inst] = true;
            }
        }
}

void DeadCodeEli::mark()
{
    while (!worklist.empty())
    {
        Instruction* inst = worklist.begin()->first;
        worklist.erase(inst);
        
        for (auto op : inst->get_operands())
        {
            Instruction* op_inst = dynamic_cast<Instruction*>(op);
            if (op_inst != nullptr)
            {
                if (marked.find(op_inst) == marked.end())
                {
                    marked[op_inst] = true;
                    worklist[op_inst] = true;
                }
            }
        }

        for (auto bb : inst->get_parent()->get_rdoms())
        {
            Instruction* inst_terminal = bb->get_terminator();
            if (marked.find(inst_terminal) == marked.end())
            {
                marked[inst_terminal] = true;
                worklist[inst_terminal] = true;
            }
        }
    }
}

void DeadCodeEli::sweep()
{
    for (auto bb : func_->get_basic_blocks())
    {
        std::vector<Instruction*> delete_list;
        for (auto inst : bb->get_instructions())
        {
            if (marked.find(inst) == marked.end())
            {
                delete_list.push_back(inst);    
            }
        }
        Instruction* terminal = nullptr;
        for (auto inst : delete_list)
        {
            if (inst == bb->get_terminator()) terminal = inst;
            else bb->delete_instr(inst);
        }
        if (terminal != nullptr)
        {
            IRStmtBuilder* builder = new IRStmtBuilder(bb, module);
            std::list<BasicBlock*> bblist;
            std::map<BasicBlock*, bool> bbcheck;
            bblist.push_back(bb);
            bbcheck[bb] = true;
            while (!bblist.empty())
            {
                auto bb_ = bblist.front();
                bblist.pop_front();
                for (auto bb_tmp : bb_->get_succ_basic_blocks())
                {
                    if (marked.find(bb_tmp->get_terminator()) != marked.end())
                    {
                        bb->delete_instr(terminal);
                        /*bb_tmp->replace_all_use_with(bb);
                        for(auto inst : bb_tmp->get_instructions())
                            bb->add_instruction(inst);
                        func_->remove(bb_tmp);*/
                        builder->create_br(bb_tmp);
                        goto Break;
                    }
                    if (bbcheck.find(bb_tmp) == bbcheck.end())
                    {
                        bbcheck[bb_tmp] = true;
                        bblist.push_back(bb_tmp);
                    }
                }
            }
Break:
            ;
        }
    }
    std::vector<BasicBlock*> useless_block_list;
    for (auto bb : func_->get_basic_blocks())
        if (bb->get_use_list().empty()) useless_block_list.push_back(bb);
    for (auto bb : useless_block_list)
        func_->remove(bb);
}

static bool IsCritical(Instruction* inst)
{
    return inst->is_ret() || inst->is_call();
}
